import os
import sys
import argparse
import matplotlib.pyplot as plt

def parse_summary(path):
    cycle = power = area = None
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line.startswith("Cycle"):
                parts = line.split()
                cycle = int(parts[2])
            elif line.startswith("Avg Power"):
                parts = line.split()
                power = float(parts[2])
            elif line.startswith("Total Area"):
                parts = line.split()
                area = float(parts[2])
    return cycle, power, area

def parse_factor_and_others(path, factor_prefix):
    factor_val = None
    other_lines = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if line.startswith(factor_prefix + ","):
                parts = line.split(",")
                try:
                    factor_val = int(parts[-1])
                except ValueError:
                    pass
            else:
                other_lines.append(line)
    if factor_val is None:
        factor_val = 0  # 如果该行不存在，默认为0
    return factor_val, tuple(other_lines)

def plot_combined(factors, cycles, powers, areas, factor_label):
    fig, ax1 = plt.subplots()

    color1 = 'blue'
    color2 = 'green'
    color3 = 'red'

    # Plot Cycles on ax1
    ln1 = ax1.plot(factors, cycles, marker='o', color=color1, label='Cycle')
    ax1.set_xlabel(f"{factor_label} (factor value)")
    ax1.set_ylabel("Cycles", color=color1)
    ax1.tick_params(axis='y', labelcolor=color1)

    # Plot Avg Power on ax2
    ax2 = ax1.twinx()
    ln2 = ax2.plot(factors, powers, marker='s', color=color2, label='Avg Power')
    ax2.set_ylabel("Avg Power (mW)", color=color2)
    ax2.tick_params(axis='y', labelcolor=color2)

    # Plot Total Area on ax3 (shifted right)
    ax3 = ax1.twinx()
    ax3.spines["right"].set_position(("axes", 1.1))
    ln3 = ax3.plot(factors, areas, marker='^', color=color3, label='Total Area')
    ax3.set_ylabel("Total Area (uM^2)", color=color3)
    ax3.tick_params(axis='y', labelcolor=color3)

    # Combine legends
    lines = ln1 + ln2 + ln3
    labels = [line.get_label() for line in lines]
    fig.legend(lines, labels, loc='center left', bbox_to_anchor=(1.05, 0.5))

    fig.tight_layout()
    plt.savefig("combined_metrics_vs_factor.png", bbox_inches='tight')
    plt.close()
    print("Saved plot as combined_metrics_vs_factor.png")

def main():
    parser = argparse.ArgumentParser(description="Plot combined metrics vs a given factor")
    parser.add_argument("dirs", nargs="+", help="Directories containing stencil_summary and config_example")
    parser.add_argument("--factor", required=True, help="Factor prefix to search")
    args = parser.parse_args()

    baseline_others = None
    data = []
    for d in args.dirs:
        summary_path = os.path.join(d, "stencil_summary")
        config_path = os.path.join(d, "config_example")
        if not os.path.isfile(summary_path) or not os.path.isfile(config_path):
            print(f"Skipping {d}: missing stencil_summary or config_example")
            continue
        factor_val, other_lines = parse_factor_and_others(config_path, args.factor)
        if baseline_others is None:
            baseline_others = other_lines
        elif other_lines != baseline_others:
            print(f"Error: config_example in '{d}' differs in non-factor lines.")
            # sys.exit(1)
        cycle, power, area = parse_summary(summary_path)
        data.append((factor_val, cycle, power, area))

    if not data:
        print("No valid data found.")
        return

    data.sort(key=lambda x: x[0])
    factors, cycles, powers, areas = zip(*data)
    plot_combined(factors, cycles, powers, areas, args.factor)

if __name__ == "__main__":
    main()
